import os

from openai import OpenAI
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())

client = OpenAI()


def get_chat_completion(
        model,
        messages,
        seed=42,
        temperature=0,
        logprobs=True,
        top_logprobs=0,
        n=1,
        **kwargs,
):
    completion = client.chat.completions.create(
        model=model,
        messages=messages,
        seed=seed,
        temperature=temperature,
        logprobs=logprobs,
        top_logprobs=top_logprobs,
        n=n,
        **kwargs,
    )
    return completion


def get_completion(
        model,
        prompt,
        seed=42,
        temperature=0,
        logprobs=True,
        top_logprobs=0,
        n=1,
        **kwargs,
):
    messages = [
        {"role": "user", "content": prompt},
    ]
    return get_chat_completion(model, messages, seed, temperature, logprobs, top_logprobs, n, **kwargs)


def get_instruct_completion(
        model,
        prompt,
        seed=42,
        temperature=0,
        max_tokens=1024,
        **kwargs,
):
    completion = client.completions.create(
        model=model,
        prompt=prompt,
        seed=seed,
        temperature=temperature,
        max_tokens=max_tokens,
        **kwargs,
    )
    return completion


def get_content(completion, index=0):
    if not hasattr(completion.choices[index],"message"):
        return completion.choices[index].text
    return completion.choices[index].message.content


def get_logprobs(completion, index=0):
    ret = []
    for item in completion.choices[index].logprobs.content:
        ret.append((item.token, item.logprob))
    return ret


def get_system_fingerprint(completion):
    return completion.system_fingerprint


if __name__ == "__main__":
    model = "gpt-3.5-turbo"
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Say a 3 word sentence to greet me."},
    ]
    completion = get_chat_completion(model, messages, temperature=0.8, n=3)
    for i in range(3):
        print(get_content(completion, i))
        print(get_logprobs(completion, i))
    print(get_system_fingerprint(completion))
